mod handle;
mod session;

pub use session::WebSocketSession;
pub use handle::WebSocketHandler;

use std::sync::Arc;
use tokio::{net::{tcp::{OwnedReadHalf, OwnedWriteHalf}, TcpStream}, sync::RwLock};
use tracing::{error, info};
use crate::websockt::{WebSocketFrame, WebSocketPayload, WebSocketReadStream, WebSocketRole, WebSocketWriteStream};
use super::{error::ClientError, Client};

pub struct WebSocketImpl {
    write_stream: Arc<RwLock<WebSocketWriteStream<OwnedWriteHalf>>>,
    read_stream:  Arc<RwLock<WebSocketReadStream<OwnedReadHalf>>>,
}

impl WebSocketImpl {
    pub fn new(stream: TcpStream) -> Self {
        let (reader, writer) = stream.into_split();
        
        let read_stream = Arc::new(RwLock::new(
            WebSocketReadStream::new(reader)
        ));
        let write_stream = Arc::new(RwLock::new(
            WebSocketWriteStream::new(writer, WebSocketRole::Client)
        ));

        Self { 
            write_stream, read_stream
        }
    }
}

impl Client<WebSocketImpl> {
    pub async fn send_text<S: Into<String>>(&mut self, text: S) -> Result<(), ClientError> {
        self.implement.write_stream.write().await.send_text(text).await?;
        Ok(())
    }

    pub async fn close(&mut self) -> Result<(), ClientError> {
        self.implement.write_stream.write().await.send_close().await?;
        // TODO: Wating for server close
        Ok(())
    }

    pub fn start_handle<H>(&self, handler: H) 
    where 
        H: WebSocketHandler + Sync + Send + 'static
    {
        let handler = Arc::new(handler);
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<WebSocketFrame>();
        
        // Read task: read any from server, and call handle
        let read_stream_clone = self.implement.read_stream.clone();
        tokio::spawn(async move {
            let read_stream = read_stream_clone;
            let sender = WebSocketSession::new(tx);
            let handler = handler;

            handler.on_open(&sender).await;

            loop {
                let frame = read_stream.write().await.read_frame().await?;
                let frame = match frame {
                    Some(f) => f,
                    None => {
                        info!("Connection closed by client");
                        break;
                    },
                };

                match &frame.payload {
                    WebSocketPayload::Text(msg) => {
                        handler.on_message(&sender, msg).await;
                    }
                    WebSocketPayload::Close => {
                        info!("🔌 WebSocket connection closing");
                        handler.on_close_with_back(&sender).await;
                    }
                    WebSocketPayload::Ping => {
                        sender.send_frame_to_server(&WebSocketFrame::client_pong());
                    }
                    WebSocketPayload::Pong => {}
                    _ => error!("Unsupported WebSocket opcode: {:?}", frame.opcode),
                }
            }

            Ok::<(), ClientError>(())
        });
    
        let write_stream_clone = self.implement.write_stream.clone();
        tokio::spawn(async move {
            let write_stream = write_stream_clone;
            while let Some(frame) = rx.recv().await {
                write_stream.write().await.send_frame(&frame).await?;
            }
            Ok::<(), ClientError>(())
        });
    }
}